{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "796032b9",
   "metadata": {},
   "source": [
    "# 统一 Caffe 网络结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb4cfe1e",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from tvm_book.config import env\n",
    "# 设置 caffeprotobuf环境\n",
    "env.set_caffeproto(Path(env.__file__).parents[3]/\"tests/caffeproto\")\n",
    "# 设置tvm环境\n",
    "env.set_tvm(\"/media/pc/data/board/arria10/lxw/tasks/tvm-test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "91d6046b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from google.protobuf import text_format\n",
    "import caffe_pb2 as pb2\n",
    "\n",
    "temp_dir = Path(\".temp\")\n",
    "temp_dir.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48144bcf",
   "metadata": {},
   "source": [
    "## 处理 in-place 层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "55cedeb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from caffe_utils import _rebuild_layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c129556e",
   "metadata": {},
   "source": [
    "该函数用于处理 Caffe 网络中的 in-place 层，解决因输入输出绑定（in-place）导致的名称冲突问题。\n",
    " 具体来说，当某层的输入（`bottom`）和输出（`top`）名称相同时（即 in-place 层），需要将其 `top` 改为层名称（`pl.name`），并通过映射表 `changed_top_dict` 动态更新其他层的 `bottom`，确保后续层引用正确。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e49c8e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "name: \"MyNetwork\"\n",
    "layer {\n",
    "  name: \"input_layer\"\n",
    "  type: \"Input\"\n",
    "  top: \"input\"\n",
    "  input_param { shape { dim: 1 dim: 3 dim: 224 dim: 224 } }\n",
    "}\n",
    "\n",
    "layer {\n",
    "  name: \"conv1\"\n",
    "  type: \"Convolution\"\n",
    "  bottom: \"input\"\n",
    "  top: \"conv1_out\"\n",
    "}\n",
    "\n",
    "layer {\n",
    "  name: \"relu1\"\n",
    "  type: \"ReLU\"\n",
    "  bottom: \"conv1_out\"\n",
    "  top: \"conv1_out\"  # 这是 in-place 操作，输入输出名称相同\n",
    "}\n",
    "\n",
    "layer {\n",
    "  name: \"pool1\"\n",
    "  type: \"Pooling\"\n",
    "  bottom: \"conv1_out\"  # 依赖原始名称 \"conv1_out\"\n",
    "  top: \"pool1_out\"\n",
    "}\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "275251ce",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"MyNetwork\"\n",
       "layer {\n",
       "  name: \"input_layer\"\n",
       "  type: \"Input\"\n",
       "  top: \"input\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 224\n",
       "      dim: 224\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"conv1\"\n",
       "  type: \"Convolution\"\n",
       "  bottom: \"input\"\n",
       "  top: \"conv1_out\"\n",
       "}\n",
       "layer {\n",
       "  name: \"relu1\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"conv1_out\"\n",
       "  top: \"relu1\"\n",
       "}\n",
       "layer {\n",
       "  name: \"pool1\"\n",
       "  type: \"Pooling\"\n",
       "  bottom: \"relu1\"\n",
       "  top: \"pool1_out\"\n",
       "}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "_rebuild_layers(predict_net.layer)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "efdcbd90",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"MyNetwork\"\n",
       "layer {\n",
       "  name: \"input_layer\"\n",
       "  type: \"Input\"\n",
       "  top: \"input_layer\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 224\n",
       "      dim: 224\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"conv1\"\n",
       "  type: \"Convolution\"\n",
       "  bottom: \"input_layer\"\n",
       "  top: \"conv1\"\n",
       "}\n",
       "layer {\n",
       "  name: \"relu1\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"conv1\"\n",
       "  top: \"relu1\"\n",
       "}\n",
       "layer {\n",
       "  name: \"pool1\"\n",
       "  type: \"Pooling\"\n",
       "  bottom: \"relu1\"\n",
       "  top: \"pool1\"\n",
       "}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from caffe_utils import unity_struct\n",
    "\n",
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "predict_net = unity_struct(predict_net)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51c6ec5f",
   "metadata": {},
   "source": [
    "```{admonition} 为什么需要这样的处理？\n",
    "- 避免名称冲突：原 in-place 层的输入输出绑定会导致后续层无法正确引用输出名称（如 pool1 无法找到 \"conv1_out\"）。\n",
    "- 兼容性需求：部分框架（如 PyTorch、TensorFlow）不支持 Caffe 的原生 in-place 机制，需显式修改名称以保证结构清晰。\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91756cb8",
   "metadata": {},
   "source": [
    "## 统一输入节点表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "61b05395",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"单输入\"\n",
       "layer {\n",
       "  name: \"data\"\n",
       "  type: \"Input\"\n",
       "  top: \"data\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 32\n",
       "      dim: 32\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"relu\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"data\"\n",
       "  top: \"relu\"\n",
       "}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from caffe_utils import unity_inputs\n",
    "\n",
    "text = \"\"\"\n",
    "name: \"单输入\"\n",
    "# 定义所有输入名称\n",
    "input: \"data\"\n",
    "\n",
    "# 输入的维度配置\n",
    "input_dim: 1\n",
    "input_dim: 3\n",
    "input_dim: 32\n",
    "input_dim: 32\n",
    "\n",
    "layer {\n",
    "    name: \"relu\"\n",
    "\tbottom: \"data\"\n",
    "\ttop: \"relu\"\n",
    "\ttype: \"ReLU\"\n",
    "}\n",
    "\"\"\".strip()\n",
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "predict_net = unity_inputs(predict_net)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5c6294ed",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"多输入\"\n",
       "layer {\n",
       "  name: \"data1\"\n",
       "  type: \"Input\"\n",
       "  top: \"data1\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 224\n",
       "      dim: 224\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"data2\"\n",
       "  type: \"Input\"\n",
       "  top: \"data2\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 1\n",
       "      dim: 112\n",
       "      dim: 112\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"relu\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"data1\"\n",
       "  top: \"relu\"\n",
       "}\n",
       "layer {\n",
       "  name: \"add\"\n",
       "  type: \"Add\"\n",
       "  bottom: \"data2\"\n",
       "  bottom: \"relu\"\n",
       "  top: \"output\"\n",
       "}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"\"\"\n",
    "name: \"多输入\"\n",
    "# 定义所有输入名称\n",
    "input: [\"data1\", \"data2\"]\n",
    "\n",
    "# 第一个输入的维度配置\n",
    "input_shape {\n",
    "  dim: 1    # 批次大小\n",
    "  dim: 3    # 通道数\n",
    "  dim: 224  # 高度\n",
    "  dim: 224  # 宽度\n",
    "}\n",
    "\n",
    "# 第二个输入的维度配置\n",
    "input_shape {\n",
    "  dim: 1    # 批次大小\n",
    "  dim: 1    # 通道数\n",
    "  dim: 112  # 高度\n",
    "  dim: 112  # 宽度\n",
    "}\n",
    "\n",
    "layer {\n",
    "\tbottom: \"data1\"\n",
    "\ttop: \"relu\"\n",
    "\tname: \"relu\"\n",
    "\ttype: \"ReLU\"\n",
    "}\n",
    "\n",
    "layer {\n",
    "\tbottom: \"data2\"\n",
    "    bottom: \"relu\"\n",
    "\ttop: \"output\"\n",
    "\tname: \"add\"\n",
    "\ttype: \"Add\"\n",
    "}\n",
    "\"\"\".strip()\n",
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "predict_net = unity_inputs(predict_net)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aca10d7d",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"多输入\"\n",
       "layer {\n",
       "  name: \"data1\"\n",
       "  type: \"Input\"\n",
       "  top: \"data1\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 224\n",
       "      dim: 224\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"data2\"\n",
       "  type: \"Input\"\n",
       "  top: \"data2\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 1\n",
       "      dim: 112\n",
       "      dim: 112\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"relu\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"data1\"\n",
       "  top: \"relu\"\n",
       "}\n",
       "layer {\n",
       "  name: \"add\"\n",
       "  type: \"Add\"\n",
       "  bottom: \"data2\"\n",
       "  bottom: \"relu\"\n",
       "  top: \"add\"\n",
       "}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from caffe_utils import unity_struct\n",
    "\n",
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "predict_net = unity_struct(predict_net)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06169384",
   "metadata": {},
   "source": [
    "## 替换数字名称"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a065e679",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"替换数字名称\"\n",
       "layer {\n",
       "  name: \"data\"\n",
       "  type: \"Input\"\n",
       "  top: \"data\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 224\n",
       "      dim: 224\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"Relu_1\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"data\"\n",
       "  top: \"Relu_1\"\n",
       "}\n",
       "layer {\n",
       "  name: \"PRelu_2\"\n",
       "  type: \"PReLU\"\n",
       "  bottom: \"Relu_1\"\n",
       "  top: \"PRelu_2\"\n",
       "}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from caffe_utils import convert_num_to_name\n",
    "text = \"\"\"\n",
    "name: \"替换数字名称\"\n",
    "layer {\n",
    "  name: \"data\"\n",
    "  type: \"Input\"\n",
    "  top: \"181\"\n",
    "  input_param {\n",
    "    shape {\n",
    "      dim: 1\n",
    "      dim: 3\n",
    "      dim: 224\n",
    "      dim: 224\n",
    "    }\n",
    "  }\n",
    "}\n",
    "layer {\n",
    "  name: \"Relu_1\"\n",
    "  type: \"ReLU\"\n",
    "  bottom: \"181\"\n",
    "  top: \"182\"\n",
    "}\n",
    "\n",
    "layer {\n",
    "  name: \"PRelu_2\"\n",
    "  type: \"PReLU\"\n",
    "  bottom: \"182\"\n",
    "  top: \"183\"\n",
    "}\n",
    "\"\"\".strip()\n",
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "predict_net = convert_num_to_name(predict_net)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e613cdea",
   "metadata": {},
   "source": [
    "## 统一 caffe 结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9b7d13e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"替换数字名称\"\n",
       "layer {\n",
       "  name: \"data\"\n",
       "  type: \"Input\"\n",
       "  top: \"data\"\n",
       "  input_param {\n",
       "    shape {\n",
       "      dim: 1\n",
       "      dim: 3\n",
       "      dim: 224\n",
       "      dim: 224\n",
       "    }\n",
       "  }\n",
       "}\n",
       "layer {\n",
       "  name: \"Relu_1\"\n",
       "  type: \"ReLU\"\n",
       "  bottom: \"data\"\n",
       "  top: \"Relu_1\"\n",
       "}\n",
       "layer {\n",
       "  name: \"PRelu_2\"\n",
       "  type: \"PReLU\"\n",
       "  bottom: \"Relu_1\"\n",
       "  top: \"PRelu_2\"\n",
       "}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from caffe_utils import unity_struct\n",
    "\n",
    "predict_net = text_format.Merge(text, pb2.NetParameter())\n",
    "predict_net = unity_struct(predict_net)\n",
    "predict_net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4053b575",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
